package api

import (
	"fmt"
	"wgmat/pkg/core"

	"github.com/cogentcore/webgpu/wgpu"
)

// MatrixAddCommand 实现Command接口的矩阵加法命令
type MatrixAddCommand struct {
	ctx  *core.Context
	a    core.Buffer
	b    core.Buffer
	c    core.Buffer
	rows uint32
	cols uint32
}

// Encode 将命令编码到CommandEncoder
func (cmd *MatrixAddCommand) Encode(encoder *wgpu.CommandEncoder) (*wgpu.BindGroup, error) {
	// 创建uniform buffer
	params := []uint32{cmd.rows, cmd.cols, cmd.rows * cmd.cols, 0}
	paramBytes := wgpu.ToBytes(params)

	device := cmd.ctx.Device()
	paramsBuffer, err := device.CreateBufferInit(&wgpu.BufferInitDescriptor{
		Label:    "uniform-params",
		Contents: paramBytes,
		Usage:    wgpu.BufferUsageUniform | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("uniform buffer创建失败: %v", err)
	}

	// 获取kernel（复用缓存）
	kernel, err := getOrCreateAddKernel(cmd.ctx)
	if err != nil {
		paramsBuffer.Release()
		return nil, err
	}

	// 创建bind group
	bindGroupEntries := []wgpu.BindGroupEntry{
		{Binding: 0, Buffer: cmd.a.Raw(), Size: cmd.a.Size()},
		{Binding: 1, Buffer: cmd.b.Raw(), Size: cmd.b.Size()},
		{Binding: 2, Buffer: cmd.c.Raw(), Size: cmd.c.Size()},
		{Binding: 3, Buffer: paramsBuffer, Size: uint64(len(paramBytes))},
	}

	bindGroup, err := device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Layout:  kernel.Layout,
		Entries: bindGroupEntries,
	})
	if err != nil {
		paramsBuffer.Release()
		return nil, fmt.Errorf("创建bind group失败: %v", err)
	}

	// 编码compute pass
	computePass := encoder.BeginComputePass(nil)
	computePass.SetPipeline(kernel.Pipeline)
	computePass.SetBindGroup(0, bindGroup, nil)
	computePass.DispatchWorkgroups((cmd.cols+15)/16, (cmd.rows+15)/16, 1)
	computePass.End()

	return bindGroup, nil
}

// MatrixMulCommand 矩阵乘法命令（演示可扩展性）
type MatrixMulCommand struct {
	ctx   *core.Context
	a     core.Buffer
	b     core.Buffer
	c     core.Buffer
	rowsA uint32
	colsA uint32
	colsB uint32
}

// Encode 实现Command接口
func (cmd *MatrixMulCommand) Encode(encoder *wgpu.CommandEncoder) (*wgpu.BindGroup, error) {
	// 创建uniform buffer
	params := []uint32{cmd.rowsA, cmd.colsA, cmd.colsB, 0}
	paramBytes := wgpu.ToBytes(params)

	device := cmd.ctx.Device()
	paramsBuffer, err := device.CreateBufferInit(&wgpu.BufferInitDescriptor{
		Label:    "uniform-params-mul",
		Contents: paramBytes,
		Usage:    wgpu.BufferUsageUniform | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("uniform buffer创建失败: %v", err)
	}

	// 获取乘法kernel（使用分块优化版本）
	kernel, err := getOrCreateMulKernel(cmd.ctx)
	if err != nil {
		paramsBuffer.Release()
		return nil, err
	}

	// 创建bind group
	bindGroupEntries := []wgpu.BindGroupEntry{
		{Binding: 0, Buffer: cmd.a.Raw(), Size: cmd.a.Size()},
		{Binding: 1, Buffer: cmd.b.Raw(), Size: cmd.b.Size()},
		{Binding: 2, Buffer: cmd.c.Raw(), Size: cmd.c.Size()},
		{Binding: 3, Buffer: paramsBuffer, Size: uint64(len(paramBytes))},
	}

	bindGroup, err := device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Layout:  kernel.Layout,
		Entries: bindGroupEntries,
	})
	if err != nil {
		paramsBuffer.Release()
		return nil, fmt.Errorf("创建bind group失败: %v", err)
	}

	// 编码compute pass（分块计算优化版）
	// 每个workgroup处理16x16的分块，workgroup数量为ceil(维度/16)
	computePass := encoder.BeginComputePass(nil)
	computePass.SetPipeline(kernel.Pipeline)
	computePass.SetBindGroup(0, bindGroup, nil)

	// workgroup数量 = (输出矩阵维度 + TILE_SIZE - 1) / TILE_SIZE
	workgroupsX := (cmd.colsB + 15) / 16 // X维度：输出矩阵的列数（分块数量）
	workgroupsY := (cmd.rowsA + 15) / 16 // Y维度：输出矩阵的行数（分块数量）
	computePass.DispatchWorkgroups(workgroupsX, workgroupsY, 1)

	computePass.End()

	return bindGroup, nil
}

type MatrixTiledMulCommand struct {
	ctx   *core.Context
	a     core.Buffer
	b     core.Buffer
	c     core.Buffer
	rowsA uint32
	colsA uint32
	colsB uint32
	tileX int
	tileY int
}

func (cmd *MatrixTiledMulCommand) Encode(encoder *wgpu.CommandEncoder) (*wgpu.BindGroup, error) {
	params := []uint32{cmd.rowsA, cmd.colsA, cmd.colsB, 0}
	paramBytes := wgpu.ToBytes(params)

	device := cmd.ctx.Device()
	paramsBuffer, err := device.CreateBufferInit(&wgpu.BufferInitDescriptor{
		Label:    "uniform-params-tiled-mul",
		Contents: paramBytes,
		Usage:    wgpu.BufferUsageUniform | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("tiled mul uniform buffer创建失败: %v", err)
	}

	kernel, err := getOrCreateTiledMulKernel(cmd.ctx, cmd.tileX, cmd.tileY)
	if err != nil {
		paramsBuffer.Release()
		return nil, err
	}

	bindGroupEntries := []wgpu.BindGroupEntry{
		{Binding: 0, Buffer: cmd.a.Raw(), Size: cmd.a.Size()},
		{Binding: 1, Buffer: cmd.b.Raw(), Size: cmd.b.Size()},
		{Binding: 2, Buffer: cmd.c.Raw(), Size: cmd.c.Size()},
		{Binding: 3, Buffer: paramsBuffer, Size: uint64(len(paramBytes))},
	}

	bindGroup, err := device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Layout:  kernel.Layout,
		Entries: bindGroupEntries,
	})
	if err != nil {
		paramsBuffer.Release()
		return nil, fmt.Errorf("tiled mul bind group创建失败: %v", err)
	}

	// 计算 workgroup 数量（向上取整）
	groupsX := (cmd.colsB + uint32(cmd.tileX) - 1) / uint32(cmd.tileX)
	groupsY := (cmd.rowsA + uint32(cmd.tileY) - 1) / uint32(cmd.tileY)

	computePass := encoder.BeginComputePass(nil)
	computePass.SetPipeline(kernel.Pipeline)
	computePass.SetBindGroup(0, bindGroup, nil)
	computePass.DispatchWorkgroups(groupsX, groupsY, 1)
	computePass.End()

	return bindGroup, nil
}

// MatrixQRCommand 实现Command接口的QR分解命令
// 使用Householder变换算法，支持经济模式(economy mode)和完整模式(full mode)
type MatrixQRCommand struct {
	ctx     *core.Context
	q       core.Buffer
	r       core.Buffer
	m, n, k uint32
}

// Encode 将QR分解命令编码到CommandEncoder
func (cmd *MatrixQRCommand) Encode(encoder *wgpu.CommandEncoder) (*wgpu.BindGroup, error) {
	device := cmd.ctx.Device()

	// 创建 uniform buffer（固定 16 字节）
	paramsBuf, err := device.CreateBuffer(&wgpu.BufferDescriptor{
		Label: "qr-params",
		Size:  16, // 4 * uint32
		Usage: wgpu.BufferUsageUniform | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("uniform buffer创建失败: %v", err)
	}
	// ✅ 添加defer确保错误时释放
	defer func() {
		if err != nil {
			paramsBuf.Release()
		}
	}()

	// 创建临时 buffer 存储 Householder 向量 v 和 vTv
	vBufferSize := (cmd.m + 1) * 4 // m 个 f32 存 v，1 个 f32 存 vTv
	vBuffer, err := device.CreateBuffer(&wgpu.BufferDescriptor{
		Label: "qr-temp-v",
		Size:  uint64(vBufferSize),
		Usage: wgpu.BufferUsageStorage,
	})
	if err != nil {
		return nil, fmt.Errorf("临时 v buffer 创建失败: %v", err)
	}
	// ✅ 添加defer确保错误时释放
	defer func() {
		if err != nil {
			vBuffer.Release()
		}
	}()

	// 获取 QR kernel
	kernel, err := getOrCreateQRKernel(cmd.ctx)
	if err != nil {
		return nil, err
	}

	// 创建 bind group
	bindGroup, err := device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Layout: kernel.Layout,
		Entries: []wgpu.BindGroupEntry{
			{Binding: 0, Buffer: cmd.q.Raw(), Size: cmd.q.Size()},
			{Binding: 1, Buffer: cmd.r.Raw(), Size: cmd.r.Size()},
			{Binding: 2, Buffer: paramsBuf, Size: 16},
			{Binding: 3, Buffer: vBuffer, Size: uint64(vBufferSize)},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("创建 bind group 失败: %v", err)
	}

	// 对每一列执行一次 kernel
	for col := uint32(0); col < cmd.k; col++ {
		params := []uint32{cmd.m, cmd.n, cmd.k, col}
		paramBytes := wgpu.ToBytes(params)
		device.GetQueue().WriteBuffer(paramsBuf, 0, paramBytes)

		computePass := encoder.BeginComputePass(nil)
		computePass.SetPipeline(kernel.Pipeline)
		computePass.SetBindGroup(0, bindGroup, nil)

		// ✅ 每次只 dispatch 一个 workgroup（256 线程），适用于 m <= 256
		// 若 m > 256，内核中的 for 循环仍能处理
		workgroups := (cmd.m + 255) / 256
		computePass.DispatchWorkgroups(workgroups, 1, 1)
		computePass.End()
	}

	// ✅ 明确buffer生命周期：paramsBuf 和 vBuffer 由返回的BindGroup管理
	// Executor会在执行后自动释放BindGroup，从而释放这两个buffer
	return bindGroup, nil
}
